import collections
from datetime import datetime
from neuro.exceptions import BaseException
from neuro.exceptions import IllegalArgumentException

def matrix_wais3_score(data):
	'''
	Calculate Matrix Reasoning WAIS-III score from iCognition CSV. Tested against iCognition data vGSP1p0 - vGSP1p8::
	
		import csv
		import StringIO
		from neuro.apps.xnat import Xnat
		from neuro.contrib.icog import matrix_wais3_score
				
		xnat = Xnat.create("gspcentral")
		details = xnat.getAssessorFilesDetails('080101_AB1234C', assessor_name='CogAssess1', selectors={'Name': 'AB1234C.csv'})
		content = xnat.getPath(details[0]['URI'])
		reader = csv.reader(StringIO.StringIO(content), delimiter=',')
		score = matrix_wais3_score(reader)	
	
	:param data: iCognition data iterable
	:type data: Iterable, csv.reader
	:returns: Matrix Reasoning WAIS score
	:rtype: float
	:author: aholmes (originally from MatrixFast.py)
	'''
	if not isinstance(data, collections.Iterable):
		raise IllegalArgumentException("Input data must be iterable")
	
	## --- drop unnecessary rows
	array = []
	for row in data:
		if(row[1] == "MatrixReason"):
			array.append([row[1], row[2], row[3], row[7],row[8],row[9]])
	
	## --- null/omitted value
	omitted = "99"

	'''
	Scoring rules
	
	1. Examinee receives 1 point for each correct response. 
	2. If the examinee obtains perfect scores on items 4 and 5, give full credit for items 1-3. 
	3. Discontinue after 4 consecutive errors or 4 errors on five consecutive trials. 
	4. Count trials with RT < 300ms as errors.
	'''
	## --- initialize scoring variables
 	Matrix_OmissionCount = 0
	Matrix_IncludedCount = 0
	Matrix_CorrectResp = 0
	Matrix_final_score = 0
	Matrix_runerrorCount = []

	## --- process data
	quitter = False
	for row in array:
		if "MatrixReason" in row[0]:
			## --- check if subject quit the task
			if row[4] == "":
				quitter = True
				break
			if row[4] == omitted:
				Matrix_OmissionCount += 1
				Matrix_runerrorCount.append(1)
			elif row[4] != omitted: 
				## --- count trials where RT < 300ms as errors
				if int(row[5]) < 300:
					Matrix_OmissionCount += 1
					Matrix_runerrorCount.append(1)
					## --- if trials 4 and 5 are correct give full credit for trials 1-3
					if int(row[1]) == 5:
						recentTrial = int(row[1])
						if sum(Matrix_runerrorCount[recentTrial-2:recentTrial]) == 0:
							Matrix_CorrectResp = 5
					## --- stop scoring if more than 4 consecutive errors
					if int(row[1]) >= 4:
						recentTrial = (int(row[1]))
						if sum(Matrix_runerrorCount[recentTrial-4:recentTrial]) >= 4:
							Matrix_final_score = Matrix_CorrectResp
							break
					## --- stop scoring if more that 4 errors the previous 5 trials
					if int(row[1]) >= 5:
						recentTrial = int(row[1])
						if sum(Matrix_runerrorCount[recentTrial-5:recentTrial]) >= 4:
							Matrix_final_score = Matrix_CorrectResp
							break
				elif int(row[5]) >= 300:
					Matrix_IncludedCount += 1	
					if int(row[1]) <= 5:
						recentTrial = int(row[1])
						## --- count correct responses
						if row[3] == row[4]:		
							Matrix_CorrectResp += 1
							Matrix_runerrorCount.append(0)
						## --- count error reponses
						elif row[4] != row[3]:			
							Matrix_runerrorCount.append(1)
					## --- if trials 4 and 5 are correct give full credit for trials 1-3
					if int(row[1]) == 5:
						recentTrial = int(row[1])
						if sum(Matrix_runerrorCount[recentTrial-2:recentTrial]) == 0:
							Matrix_CorrectResp = 5
					## --- stop scoring if more than 4 consecutive errors
					if int(row[1]) >= 4:
						recentTrial = int(row[1])
						if sum(Matrix_runerrorCount[recentTrial-4:recentTrial]) >= 4:
							Matrix_final_score = Matrix_CorrectResp
							break
					if int(row[1]) > 5:
						## --- count correct responses
						if row[3] == row[4]:			
							Matrix_CorrectResp += 1
							Matrix_runerrorCount.append(0)
						## --- count error reponses
						elif row[4] != row[3]:	
							Matrix_runerrorCount.append(1)	
					## --- stop scoring if more that 4 errors the previous 5 trials
					if int(row[1]) >= 5:
						recentTrial = int(row[1])
						if sum(Matrix_runerrorCount[recentTrial-5:recentTrial]) >= 4:
							Matrix_final_score = Matrix_CorrectResp
							break
						Matrix_final_score = Matrix_CorrectResp
	
	if(quitter):
		return None
	
	return Matrix_final_score

def get_task_data(data, task_name=None):
	'''
	Pull out specific task-related data from iCognition CSV. Tested against iCognition data vGSP1p0 - vGSP1p8::
	
		import csv
		import StringIO
		from neuro.apps.xnat import Xnat
		from neuro.contrib.icog import get_task_data
		
		xnat = Xnat.create("gspcentral")
		details = xnat.getAssessorFilesDetails('GSPCentral_S00001_MR_1', assessor_name='CogAssess1', selectors={'Name': 'ABCDE.csv'})
		content = xnat.getPath(details[0]['URI'])
		reader = csv.reader(StringIO.StringIO(content), delimiter=',')
		data = get_task_data(reader, task_name="MatrixReason")

		for row in data:
			print "Task=" + row["TASK"] + ", Trial=" + row["TRIAL"]
	
	:param data: iCognition data as Iterable e.g., csv.reader
	:type data: Iterable, csv.reader
	:param task_name: Task name e.g., 'MatrixReason'
	:type task_name: str
	:returns: Task-related data as a list of dicts
	:rtype: list
	:author: tokeefe
	'''
	if not isinstance(data, collections.Iterable):
		raise IllegalArgumentException("Input data must be iterable")
	
	## --- get headers
	headers = data.next()
	
	## --- get 'TASK' column index
	task_ind = headers.index("TASK")
	
	## --- get task specific data
	task_data = []
	for row in data:
		if task_name == None or row[task_ind] == task_name:
			## --- zip labels with rows and add to task_data
			labeled_row = dict(zip(headers, row)) 
			task_data.append(labeled_row)
	
	return task_data

def compliance_score(data):
	'''
	Calculate iCognition compliance score::
	
		import csv
		import StringIO
		from neuro.apps.xnat import Xnat
		from neuro.contrib.icog import matrix_wais3_score
				
		xnat = Xnat.create("gspcentral")
		details = xnat.getAssessorFilesDetails('080101_AB1234C', assessor_name='CogAssess1', selectors={'Name': 'AB1234C.csv'})
		content = xnat.getPath(details[0]['URI'])
		reader = csv.reader(StringIO.StringIO(content), delimiter=',')
		score = compliance_score(reader)

	:param data: iCognition data 
	:type data: Iterable, csv.reader
	:returns: Compliance score
	:rtype: float
	:author: aholmes
	'''
	iCOG_Comp = 0
	flanker3sec = False
	flanker5sec = False
	stressRatings = []
	performRatings = []
	attendRatings = []
	for row in data:
		## --- -2 Comp for taking under 25 minutes
		if "cog_starttime" in row[4]:
			T1 = datetime.strptime(row[8], "%y%m%d %H:%M:%S")
		if "cog_endtime" in row[4]:
			T2 = datetime.strptime(row[8], "%y%m%d %H:%M:%S")
			delta = T2 - T1
			if delta.seconds < 1500:
				iCOG_Comp -= 2
				
		## --- -1 Comp for each keyboarding error
		if "Keyboarding" in row[1]:
			if int(row[11]) == 0:
				iCOG_Comp -= 1
				
		## --- -1 Comp for each error in the first three Matrix Reasoning trials
		if "MatrixReason" in row[1]:
			if (int(row[2]) <= 3) and (int(row[11]) == 0):
				iCOG_Comp -= 1
				
		## --- -2 Comp if Flanker RT > 3 seconds or -4 Comp if > 5 seconds
		if "Flanker" in row[1]:
			if (int(row[9]) >= 3000):
				flanker3sec = True
				if (int(row[9]) >= 5000):
					flanker5sec = True
					
		## --- -1 Comp for error in RMET catch trials
		if "RMET" in row[1]:
			if (row[5] == "Catch") and (int(row[11]) == 0):
				iCOG_Comp -= 1

		## --- Logs the Stress/Performance/Attend ratings for each measure
		if "RateTheTask" in row[1]:
                        if (row[4] == "stress"):
				stressRatings.append(row[8])
                        if (row[4] == "perform"): 
				performRatings.append(row[8])
	                if (row[4] == "attend"):
				attendRatings.append(row[8])

	## --- -2 Comp if all Stress/Performance/Attend ratings for each measure are identical
	if len(set(stressRatings)) == 1:
		iCOG_Comp -= 2
	if len(set(performRatings)) == 1:
		iCOG_Comp -= 2
	if len(set(attendRatings)) == 1:
		iCOG_Comp -= 2

	## --- subtract flanker compliance from iCOG_Comp
	if flanker3sec and flanker5sec:
		 iCOG_Comp -= 4
	elif flanker3sec and not flanker5sec:
		iCOG_Comp -= 2

	return iCOG_Comp
